import torch
import numpy as np
import matplotlib.pyplot as plt
import cv2

# 动态设备分配
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def load_model(model_path):
    """
    加载保存为完整模型的 .pth 文件
    """
    model = torch.load(model_path, map_location=device)
    model.eval()  # 切换到评估模式
    print(f"Model successfully loaded on {device}")
    return model

def preprocess_image(image_path):
    """
    预处理3D NumPy数组，转为 5D 张量输入 (batch_size, channels, depth, height, width)
    """
    image_array = np.load(image_path)  # 加载3D图像数据
    image_array = (image_array - np.min(image_array)) / (np.max(image_array) - np.min(image_array))  # 归一化到 [0, 1]

    # 增加维度形成 (1, channels, depth, height, width)
    if len(image_array.shape) == 3:  # 如果是3D图像数据，添加通道维度和批次维度
        image_array = np.expand_dims(image_array, axis=0)  # (1, depth, height, width)
        image_array = np.expand_dims(image_array, axis=0)  # (1, 1, depth, height, width)

    # 转为 PyTorch 张量并移动到设备
    tensor = torch.tensor(image_array, dtype=torch.float32).to(device)

    # 调试：打印张量形状和数据范围
    print(f"Preprocessed tensor shape: {tensor.shape}")
    print(f"Tensor min: {tensor.min()}, max: {tensor.max()}")
    print(f"Tensor mean: {tensor.mean()}, std: {tensor.std()}")
    return tensor

def apply_highlight(input_slice, prediction_slice, threshold=0.5):
    """
    在预测结果上应用荧光效果
    input_slice: 输入图像的切片（灰度图）
    prediction_slice: 预测结果的切片
    threshold: 阈值，用于确定目标区域
    """
    # 将预测结果转换为二值图像
    _, binary_prediction = cv2.threshold(prediction_slice, threshold, 255, cv2.THRESH_BINARY)

    # 将灰度图像转换为 RGB 图像
    input_rgb = cv2.cvtColor(input_slice, cv2.COLOR_GRAY2RGB)

    # 创建一个与输入图像相同大小的全零 RGB 图像
    highlighted = np.zeros_like(input_rgb)

    # 将二值图像中的白色区域设置为绿色
    highlighted[binary_prediction == 255] = [0, 255, 0]

    # 将高亮图像与输入图像合并
    result = cv2.addWeighted(input_rgb, 0.7, highlighted, 0.3, 0)
    return result

if __name__ == "__main__":
    # 加载模型
    model_path = "./models/best_model.pth"  # 替换为实际路径
    model = load_model(model_path)

    # 加载测试数据
    image_path = "./static/img/t1.npy"  # 替换为实际3D图像数据路径
    input_tensor = preprocess_image(image_path)

    # 可视化输入图像的中间切片
    depth = input_tensor.shape[2] // 2  # 计算深度中间位置

    # 模型预测
    with torch.no_grad():  # 禁用梯度计算
        output = model(input_tensor)
        output = torch.sigmoid(output)  # 转换为概率图
        output = output[0, 1].detach().cpu().numpy()  # 提取预测结果 (3D NumPy array)

    # 将输入张量转换为 NumPy 数组
    input_array = input_tensor[0, 0].cpu().numpy()  # (depth, height, width)

    # 提取中间切片
    input_slice = input_array[depth, :, :]
    prediction_slice = output[depth, :, :]

    # 应用荧光效果
    highlighted_image = apply_highlight(input_slice, prediction_slice)

    # 将高亮图像转换为 uint8 类型并缩放到 [0, 255]
    highlighted_image = (highlighted_image * 255).astype(np.uint8)

    # 显示高亮图像
    plt.imshow(cv2.cvtColor(highlighted_image, cv2.COLOR_BGR2RGB))
    plt.title("Highlighted Prediction on Input Image")
    plt.axis('off')
    plt.show()